import sys
import os
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QPushButton, QFileDialog, QSlider, QSpinBox, QComboBox,
                             QGroupBox, QGridLayout, QAction, QMenuBar, QSplitter, QMessageBox,
                             QScrollArea, QSizePolicy, QInputDialog)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor, QCursor, QTransform
from PyQt5.QtCore import Qt, QPoint, QRect, QSize

# 导入OpenCV用于图像处理
import cv2


class ImageProcessingPlatform(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.current_image = None
        self.original_image = None
        self.temp_image = None
        self.shapes = []
        self.cropping = False
        self.crop_start = QPoint()
        self.crop_end = QPoint()
        self.images_to_stitch = []
        self.is_stitching = False

    def initUI(self):
        # 设置窗口标题和大小
        self.setWindowTitle('图像处理平台')
        self.setGeometry(100, 100, 1200, 800)

        # 创建菜单栏
        self.createMenus()

        # 创建主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)

        # 创建左侧控制面板
        control_panel = QWidget()
        control_layout = QVBoxLayout(control_panel)
        control_panel.setFixedWidth(300)

        # 图像加载组
        load_group = QGroupBox("图像加载")
        load_layout = QVBoxLayout(load_group)

        self.load_button = QPushButton("加载图像")
        self.load_button.clicked.connect(self.loadImage)

        self.stitch_button = QPushButton("添加到拼接")
        self.stitch_button.clicked.connect(self.addToStitch)
        self.stitch_button.setEnabled(False)

        self.perform_stitch_button = QPushButton("执行图像拼接")
        self.perform_stitch_button.clicked.connect(self.performStitch)
        self.perform_stitch_button.setEnabled(False)

        load_layout.addWidget(self.load_button)
        load_layout.addWidget(self.stitch_button)
        load_layout.addWidget(self.perform_stitch_button)

        # 图像编辑组
        edit_group = QGroupBox("图像编辑")
        edit_layout = QVBoxLayout(edit_group)

        # 旋转控件
        rotate_layout = QHBoxLayout()
        rotate_label = QLabel("旋转角度:")
        self.rotate_spin = QSpinBox()
        self.rotate_spin.setRange(-359, 359)
        self.rotate_button = QPushButton("旋转")
        self.rotate_button.clicked.connect(self.rotateImage)

        rotate_layout.addWidget(rotate_label)
        rotate_layout.addWidget(self.rotate_spin)
        rotate_layout.addWidget(self.rotate_button)

        # 缩放控件
        scale_layout = QHBoxLayout()
        scale_label = QLabel("缩放比例(%):")
        self.scale_spin = QSpinBox()
        self.scale_spin.setRange(10, 500)
        self.scale_spin.setValue(100)
        self.scale_button = QPushButton("缩放")
        self.scale_button.clicked.connect(self.scaleImage)

        scale_layout.addWidget(scale_label)
        scale_layout.addWidget(self.scale_spin)
        scale_layout.addWidget(self.scale_button)

        # 裁剪控件
        self.crop_button = QPushButton("裁剪图像")
        self.crop_button.clicked.connect(self.startCrop)

        edit_layout.addLayout(rotate_layout)
        edit_layout.addLayout(scale_layout)
        edit_layout.addWidget(self.crop_button)

        # 色彩调整组
        color_group = QGroupBox("色彩调整")
        color_layout = QGridLayout(color_group)

        # 亮度调整
        brightness_label = QLabel("亮度:")
        self.brightness_slider = QSlider(Qt.Horizontal)
        self.brightness_slider.setRange(-100, 100)
        self.brightness_slider.setValue(0)
        self.brightness_slider.valueChanged.connect(self.adjustColor)

        # 对比度调整
        contrast_label = QLabel("对比度:")
        self.contrast_slider = QSlider(Qt.Horizontal)
        self.contrast_slider.setRange(-100, 100)
        self.contrast_slider.setValue(0)
        self.contrast_slider.valueChanged.connect(self.adjustColor)

        # 饱和度调整
        saturation_label = QLabel("饱和度:")
        self.saturation_slider = QSlider(Qt.Horizontal)
        self.saturation_slider.setRange(-100, 100)
        self.saturation_slider.setValue(0)
        self.saturation_slider.valueChanged.connect(self.adjustColor)

        color_layout.addWidget(brightness_label, 0, 0)
        color_layout.addWidget(self.brightness_slider, 0, 1)
        color_layout.addWidget(contrast_label, 1, 0)
        color_layout.addWidget(self.contrast_slider, 1, 1)
        color_layout.addWidget(saturation_label, 2, 0)
        color_layout.addWidget(self.saturation_slider, 2, 1)

        # 平滑处理组
        smooth_group = QGroupBox("平滑处理")
        smooth_layout = QVBoxLayout(smooth_group)

        self.smooth_combo = QComboBox()
        self.smooth_combo.addItems(["均值滤波", "高斯滤波", "中值滤波", "双边滤波"])

        self.smooth_button = QPushButton("应用平滑")
        self.smooth_button.clicked.connect(self.applySmoothing)

        smooth_layout.addWidget(self.smooth_combo)
        smooth_layout.addWidget(self.smooth_button)

        # 形状识别组
        shape_group = QGroupBox("形状识别")
        shape_layout = QVBoxLayout(shape_group)

        self.shape_combo = QComboBox()
        self.shape_combo.addItems(["检测所有形状", "检测圆形", "检测矩形", "检测三角形"])

        self.shape_button = QPushButton("识别形状")
        self.shape_button.clicked.connect(self.detectShapes)

        shape_layout.addWidget(self.shape_combo)
        shape_layout.addWidget(self.shape_button)

        # 添加所有组到控制面板
        control_layout.addWidget(load_group)
        control_layout.addWidget(edit_group)
        control_layout.addWidget(color_group)
        control_layout.addWidget(smooth_group)
        control_layout.addWidget(shape_group)
        control_layout.addStretch()

        # 创建右侧图像显示区域
        self.scroll_area = QScrollArea()
        self.scroll_area.setWidgetResizable(True)

        self.image_label = QLabel("请加载图像")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(1, 1)  # 允许标签缩小
        self.image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.image_label.mousePressEvent = self.mousePress
        self.image_label.mouseMoveEvent = self.mouseMove
        self.image_label.mouseReleaseEvent = self.mouseRelease

        self.scroll_area.setWidget(self.image_label)

        # 添加到主布局
        main_layout.addWidget(control_panel)
        main_layout.addWidget(self.scroll_area, 1)

    def createMenus(self):
        # 创建文件菜单
        file_menu = self.menuBar().addMenu("文件")

        # 添加打开动作
        open_action = QAction("打开", self)
        open_action.setShortcut("Ctrl+O")
        open_action.triggered.connect(self.loadImage)
        file_menu.addAction(open_action)

        # 添加保存动作
        save_action = QAction("保存", self)
        save_action.setShortcut("Ctrl+S")
        save_action.triggered.connect(self.saveImage)
        file_menu.addAction(save_action)

        # 添加退出动作
        exit_action = QAction("退出", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)

        # 创建编辑菜单
        edit_menu = self.menuBar().addMenu("编辑")

        # 添加撤销动作
        undo_action = QAction("撤销", self)
        undo_action.setShortcut("Ctrl+Z")
        undo_action.triggered.connect(self.undo)
        edit_menu.addAction(undo_action)

        # 创建视图菜单
        view_menu = self.menuBar().addMenu("视图")

        # 添加适应窗口动作
        fit_action = QAction("适应窗口", self)
        fit_action.setShortcut("Ctrl+F")
        fit_action.triggered.connect(self.fitToWindow)
        view_menu.addAction(fit_action)

        # 添加实际大小动作
        actual_action = QAction("实际大小", self)
        actual_action.setShortcut("Ctrl+A")
        actual_action.triggered.connect(self.actualSize)
        view_menu.addAction(actual_action)

    def loadImage(self):
        # 打开文件对话框选择图像
        file_name, _ = QFileDialog.getOpenFileName(
            self, "选择图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp *.gif)"
        )

        if file_name:
            # 保存原始图像
            self.original_image = cv2.imread(file_name)
            self.current_image = self.original_image.copy()
            self.temp_image = self.current_image.copy()
            self.shapes = []
            self.cropping = False
            self.images_to_stitch = []
            self.stitch_button.setEnabled(True)
            self.perform_stitch_button.setEnabled(False)
            self.updateImage()

    def updateImage(self):
        if self.current_image is not None:
            # 转换OpenCV图像为Qt可显示格式
            height, width, channel = self.current_image.shape
            bytes_per_line = 3 * width
            q_img = QImage(self.current_image.data, width, height, bytes_per_line, QImage.Format_BGR888)

            # 如果有形状检测结果，绘制它们
            if self.shapes:
                painter = QPainter(q_img)
                painter.setRenderHint(QPainter.Antialiasing)

                for shape in self.shapes:
                    shape_type, contour = shape
                    if shape_type == "圆形":
                        # 绘制圆形
                        (x, y), radius = cv2.minEnclosingCircle(contour)
                        center = (int(x), int(y))
                        radius = int(radius)
                        painter.setPen(QPen(QColor(255, 0, 0), 2))
                        painter.drawEllipse(center[0] - radius, center[1] - radius, radius * 2, radius * 2)
                        painter.drawText(center[0] - 10, center[1] - radius - 10, "圆形")
                    elif shape_type == "矩形":
                        # 绘制矩形
                        x, y, w, h = cv2.boundingRect(contour)
                        painter.setPen(QPen(QColor(0, 255, 0), 2))
                        painter.drawRect(x, y, w, h)
                        painter.drawText(x, y - 10, "矩形")
                    elif shape_type == "三角形":
                        # 绘制三角形
                        epsilon = 0.02 * cv2.arcLength(contour, True)
                        approx = cv2.approxPolyDP(contour, epsilon, True)
                        if len(approx) == 3:
                            points = [(p[0][0], p[0][1]) for p in approx]
                            painter.setPen(QPen(QColor(0, 0, 255), 2))
                            for i in range(3):
                                painter.drawLine(points[i][0], points[i][1], points[(i + 1) % 3][0],
                                                 points[(i + 1) % 3][1])
                            painter.drawText(points[0][0], points[0][1] - 10, "三角形")
                    elif shape_type == "多边形":
                        # 绘制多边形
                        epsilon = 0.02 * cv2.arcLength(contour, True)
                        approx = cv2.approxPolyDP(contour, epsilon, True)
                        if len(approx) > 4:
                            points = [(p[0][0], p[0][1]) for p in approx]
                            painter.setPen(QPen(QColor(255, 255, 0), 2))
                            for i in range(len(points)):
                                painter.drawLine(points[i][0], points[i][1], points[(i + 1) % len(points)][0],
                                                 points[(i + 1) % len(points)][1])
                            painter.drawText(points[0][0], points[0][1] - 10, f"多边形({len(points)}边)")

                painter.end()

            # 如果正在裁剪，绘制裁剪框
            if self.cropping and not self.crop_start.isNull() and not self.crop_end.isNull():
                painter = QPainter(q_img)
                painter.setPen(QPen(QColor(255, 0, 0), 2, Qt.DashLine))

                # 确保坐标是正确的顺序
                x = min(self.crop_start.x(), self.crop_end.x())
                y = min(self.crop_start.y(), self.crop_end.y())
                w = abs(self.crop_start.x() - self.crop_end.x())
                h = abs(self.crop_start.y() - self.crop_end.y())

                painter.drawRect(x, y, w, h)
                painter.end()

            # 显示图像
            pixmap = QPixmap.fromImage(q_img)
            self.image_label.setPixmap(pixmap.scaled(
                self.image_label.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            ))

    def saveImage(self):
        if self.current_image is not None:
            file_name, _ = QFileDialog.getSaveFileName(
                self, "保存图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
            )

            if file_name:
                try:
                    cv2.imwrite(file_name, self.current_image)
                    QMessageBox.information(self, "成功", "图像保存成功!")
                except Exception as e:
                    QMessageBox.critical(self, "错误", f"保存图像时出错: {str(e)}")

    def undo(self):
        if self.original_image is not None:
            self.current_image = self.original_image.copy()
            self.temp_image = self.current_image.copy()
            self.shapes = []
            self.cropping = False
            self.updateImage()

    def fitToWindow(self):
        if not self.image_label.pixmap().isNull():
            self.image_label.adjustSize()

    def actualSize(self):
        if not self.image_label.pixmap().isNull():
            self.image_label.setPixmap(self.image_label.pixmap().scaled(
                self.current_image.shape[1],
                self.current_image.shape[0],
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            ))

    def rotateImage(self):
        if self.current_image is not None:
            angle = self.rotate_spin.value()
            if angle != 0:
                rows, cols = self.current_image.shape[:2]
                M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
                self.current_image = cv2.warpAffine(self.current_image, M, (cols, rows))
                self.temp_image = self.current_image.copy()
                self.updateImage()

    def scaleImage(self):
        if self.current_image is not None:
            scale_percent = self.scale_spin.value()
            width = int(self.current_image.shape[1] * scale_percent / 100)
            height = int(self.current_image.shape[0] * scale_percent / 100)
            dim = (width, height)
            self.current_image = cv2.resize(self.current_image, dim, interpolation=cv2.INTER_AREA)
            self.temp_image = self.current_image.copy()
            self.updateImage()

    def startCrop(self):
        if self.current_image is not None:
            self.cropping = True
            self.crop_start = QPoint()
            self.crop_end = QPoint()
            self.setCursor(Qt.CrossCursor)
            self.updateImage()

    def mousePress(self, event):
        if self.cropping and event.button() == Qt.LeftButton:
            self.crop_start = event.pos()
            self.crop_end = QPoint()
            self.updateImage()

    def mouseMove(self, event):
        if self.cropping and (event.buttons() & Qt.LeftButton):
            self.crop_end = event.pos()
            self.updateImage()

    def mouseRelease(self, event):
        if self.cropping and event.button() == Qt.LeftButton:
            self.crop_end = event.pos()
            self.setCursor(Qt.ArrowCursor)

            # 确保坐标是正确的顺序
            x = min(self.crop_start.x(), self.crop_end.x())
            y = min(self.crop_start.y(), self.crop_end.y())
            w = abs(self.crop_start.x() - self.crop_end.x())
            h = abs(self.crop_start.y() - self.crop_end.y())

            if w > 0 and h > 0:
                # 裁剪图像
                if not self.image_label.pixmap().isNull():
                    pixmap = self.image_label.pixmap()
                    cropped_pixmap = pixmap.copy(x, y, w, h)

                    # 转换回OpenCV格式
                    qimg = cropped_pixmap.toImage()
                    ptr = qimg.bits()
                    ptr.setsize(qimg.byteCount())
                    arr = np.array(ptr).reshape(qimg.height(), qimg.width(), 3)
                    self.current_image = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
                    self.temp_image = self.current_image.copy()
                    self.cropping = False
                    self.updateImage()

    def adjustColor(self):
        if self.current_image is not None:
            brightness = self.brightness_slider.value()
            contrast = self.contrast_slider.value()
            saturation = self.saturation_slider.value()

            # 创建临时图像进行调整
            img = self.temp_image.copy()

            # 调整亮度和对比度
            if brightness != 0 or contrast != 0:
                alpha = 1.0 + contrast / 100.0
                beta = brightness
                img = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)

            # 调整饱和度
            if saturation != 0:
                hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
                h, s, v = cv2.split(hsv)

                if saturation > 0:
                    limit = 255 - saturation
                    s[s > limit] = 255
                    s[s <= limit] += saturation
                else:
                    s = np.maximum(s, np.zeros(s.shape) - saturation)
                    s = s + saturation

                hsv = cv2.merge([h, s, v])
                img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

            self.current_image = img
            self.updateImage()

    def applySmoothing(self):
        if self.current_image is not None:
            method = self.smooth_combo.currentText()
            img = self.current_image.copy()

            if method == "均值滤波":
                img = cv2.blur(img, (5, 5))
            elif method == "高斯滤波":
                img = cv2.GaussianBlur(img, (5, 5), 0)
            elif method == "中值滤波":
                img = cv2.medianBlur(img, 5)
            elif method == "双边滤波":
                img = cv2.bilateralFilter(img, 9, 75, 75)

            self.current_image = img
            self.temp_image = self.current_image.copy()
            self.updateImage()

    def detectShapes(self):
        if self.current_image is not None:
            method = self.shape_combo.currentText()
            img = self.current_image.copy()
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            _, thresh = cv2.threshold(blurred, 127, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

            # 查找轮廓
            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            self.shapes = []

            for contour in contours:
                area = cv2.contourArea(contour)
                # 忽略小面积区域
                if area < 100:
                    continue

                if method == "检测所有形状" or method == "检测圆形":
                    # 检测圆形
                    perimeter = cv2.arcLength(contour, True)
                    approx = cv2.approxPolyDP(contour, 0.04 * perimeter, True)
                    (x, y), radius = cv2.minEnclosingCircle(contour)
                    circle_area = np.pi * radius * radius

                    # 如果面积比例接近圆形，认为是圆形
                    if abs(area - circle_area) / circle_area < 0.3:
                        self.shapes.append(("圆形", contour))
                        if method == "检测圆形":
                            continue

                if method == "检测所有形状" or method == "检测矩形":
                    # 检测矩形
                    perimeter = cv2.arcLength(contour, True)
                    approx = cv2.approxPolyDP(contour, 0.04 * perimeter, True)

                    if len(approx) == 4:
                        self.shapes.append(("矩形", contour))
                        if method == "检测矩形":
                            continue

                if method == "检测所有形状" or method == "检测三角形":
                    # 检测三角形
                    perimeter = cv2.arcLength(contour, True)
                    approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)

                    if len(approx) == 3:
                        self.shapes.append(("三角形", contour))
                        if method == "检测三角形":
                            continue

                # 如果是检测所有形状且未匹配到其他形状，认为是多边形
                if method == "检测所有形状":
                    self.shapes.append(("多边形", contour))

            self.updateImage()

    def addToStitch(self):
        if self.current_image is not None:
            self.images_to_stitch.append(self.current_image.copy())
            QMessageBox.information(self, "成功", f"已添加图像到拼接队列，当前共有 {len(self.images_to_stitch)} 张图像")
            self.perform_stitch_button.setEnabled(len(self.images_to_stitch) >= 2)

    def performStitch(self):
        if len(self.images_to_stitch) < 2:
            QMessageBox.warning(self, "警告", "至少需要两张图像才能进行拼接")
            return

        # 创建图像拼接器
        stitcher = cv2.Stitcher_create() if hasattr(cv2, 'Stitcher_create') else cv2.Stitcher_create(cv2.Stitcher_SCANS)

        # 执行拼接
        status, stitched = stitcher.stitch(self.images_to_stitch)

        if status == cv2.Stitcher_OK:
            self.current_image = stitched
            self.temp_image = self.current_image.copy()
            self.shapes = []
            self.updateImage()
            QMessageBox.information(self, "成功", "图像拼接成功!")
        else:
            QMessageBox.critical(self, "错误", f"图像拼接失败，错误码: {status}")


if __name__ == '__main__':
    # 创建Qt应用
    app = QApplication(sys.argv)

    # 确保中文显示正常
    font = app.font()
    font.setFamily("SimHei")
    app.setFont(font)

    # 创建并显示主窗口
    window = ImageProcessingPlatform()
    window.show()

    # 进入应用主循环
    sys.exit(app.exec_())